#!/usr/bin/env python3

from ete4 import Tree
from ete4.treeview import faces, AttrFace, TreeStyle, NodeStyle


def layout(node):
    if node.is_leaf:
        N = AttrFace('name', fsize=30)
        faces.add_face_to_node(N, node, 0, position='aligned')


def get_example_tree():
    # Set dashed blue lines in all leaves.
    nst1 = NodeStyle()
    nst1['bgcolor'] = 'LightSteelBlue'
    nst2 = NodeStyle()
    nst2['bgcolor'] = 'Moccasin'
    nst3 = NodeStyle()
    nst3['bgcolor'] = 'DarkSeaGreen'
    nst4 = NodeStyle()
    nst4['bgcolor'] = 'Khaki'

    t = Tree('((((a1,a2),a3), ((b1,b2),(b3,b4))), ((c1,c2),c3));')
    for n in t.traverse():
        n.dist = 0

    n1 = t.common_ancestor(['a1', 'a2', 'a3'])
    n1.set_style(nst1)
    n2 = t.common_ancestor(['b1', 'b2', 'b3', 'b4'])
    n2.set_style(nst2)
    n3 = t.common_ancestor(['c1', 'c2', 'c3'])
    n3.set_style(nst3)
    n4 = t.common_ancestor(['b3', 'b4'])
    n4.set_style(nst4)
    ts = TreeStyle()
    ts.layout_fn = layout
    ts.show_leaf_name = False

    ts.mode = 'c'
    ts.root_opening_factor = 1
    return t, ts


if __name__ == '__main__':
    t, ts = get_example_tree()
    t.show(tree_style=ts)
    # Or save it with:
    #   t.render('node_background.png', w=400, tree_style=ts)
